#include <iostream>
#include "clang/AST/AST.h"
#include "clang/AST/DeclObjC.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/ASTMatchers/ASTMatchers.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/ASTMatchers/ASTMatchFinder.h"
#include "clang/Frontend/FrontendPluginRegistry.h"

using namespace clang;
using namespace std;
using namespace llvm;
using namespace clang::ast_matchers;

namespace YTSimpleCodeCheckTool {
    
    class YTMatchHandler: public MatchFinder::MatchCallback {
    private:
        CompilerInstance &CI;
        
        // 是否使用了泛型
        bool isGenericType(const string typeStr)
        {
            if (typeStr.find("<") != string::npos &&
                typeStr.find(">") != string::npos) {
                return true;
            }
            return false;
        }

        // 是否需要使用copy
        bool isShouldUseCopy(const string typeStr)
        {
            vector<string> shouldUseCopys = {"NSString","(^)"/*block*/};
            for (size_t idx = 0; idx < shouldUseCopys.size(); idx++ ) {
                const string item = shouldUseCopys[idx];
                if (typeStr.find(item) != string::npos) {
                    return isGenericType(typeStr)?false:true;
                }
            }
            return false;
        }
        
        // 是否是用户代码
        bool isUserSourceCode(const string filename)
        {
            // 非Xcode中的源码都认为是用户源码
            if (filename.empty() ||
                filename.find("/Applications/Xcode.app/") == 0) {
                return false;
            }
            return true;
        }
        
        void handlePropertyDecl(const ObjCPropertyDecl *propertyDecl) {
            ObjCPropertyDecl::PropertyAttributeKind attrKind = propertyDecl->getPropertyAttributes();
            const string typeStr = propertyDecl->getType().getAsString();
            
            if (propertyDecl->getTypeSourceInfo()) {
                DiagnosticsEngine &diag = CI.getDiagnostics();
                if (isShouldUseCopy(typeStr)) {
                    if (!(attrKind & ObjCPropertyDecl::OBJC_PR_copy)) {
                        diag.Report(propertyDecl->getBeginLoc(), diag.getCustomDiagID(DiagnosticsEngine::Warning, "❌ ⚠️ YT WARNING %0 应该使用 copy 修饰")) << typeStr;
                    }
                }
            }
        }
       
    public:
        YTMatchHandler(CompilerInstance &CI) :CI(CI) {}
        
        void run(const MatchFinder::MatchResult &Result) {
            const ObjCPropertyDecl *propertyDecl = Result.Nodes.getNodeAs<ObjCPropertyDecl>("ObjcPropertyDecl");
            
            if (propertyDecl &&
                isUserSourceCode(CI.getSourceManager().getFilename(propertyDecl->getSourceRange().getBegin()).str())) {
                handlePropertyDecl(propertyDecl);
            }
            
        }
    };
    
    class YTASTConsumer: public ASTConsumer {
    private:
        MatchFinder matcher;
        YTMatchHandler handler;
    public:
        YTASTConsumer(CompilerInstance &CI) :handler(CI) {
            matcher.addMatcher(objcPropertyDecl().bind("ObjcPropertyDecl"), &handler);
        }
        
        void HandleTranslationUnit(ASTContext &context) {
            matcher.matchAST(context);
        }
    };
    
    class YTASTAction: public PluginASTAction {
    public:
        unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI, StringRef iFile) {
            return unique_ptr<YTASTConsumer> (new YTASTConsumer(CI));
        }
        
        bool ParseArgs(const CompilerInstance &ci, const std::vector<std::string> &args) {
            return true;
        }
    };
}

// 注册PluginASTAction
static FrontendPluginRegistry::Add<YTSimpleCodeCheckTool::YTASTAction> X("YTSimpleCodeCheckTool", "A custom static code analysis of Objc");
